# Copyright 2018 Daniel 'grindhold' Brendle
#
# This file is part of cmljitter.
#
# cmljitter is free software: you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later
# version.
#
# cmljitter is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
# PURPOSE. See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public
# License along with cmljitter.
# If not, see http://www.gnu.org/licenses/.

from xml.etree.ElementTree import ElementTree, fromstring, dump, register_namespace
import sys
import random

USAGE = """
cmljitter - add random positions to atoms in a cml file

Usage: pipe a cml file to stdin and cmljitter will output
       the same cml file with slightly randomized coordinates
       on each atom.
"""

def run():
    if len(sys.argv) > 1:
        print (USAGE)
        return
    data = sys.stdin.read()
    register_namespace('', "http://www.xml-cml.org/schema")
    root = fromstring(data)

    def float2string(flt):
        return str(flt)[:(8,9)[int(flt>0)]]

    for element in root.iter():
        if element.tag == "{http://www.xml-cml.org/schema}atom":
            element.set("x3", float2string(random.random()*40))
            element.set("y3", float2string(random.random()*40))
            element.set("z3", float2string(random.random()*40))

    dump(root)

if __name__ == "__main__":
    run()
